package info.wangyuan.agent.config;

import cn.hutool.core.collection.CollUtil;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import info.wangyuan.agent.common.Constant.RedisConstant;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * Redis 对话记忆仓储（基于 Redis List，每次追加消息 + 仅保留最近 15 条）
 * 每个 memoryId -> 一个 Redis List
 * - RPUSH 追加单条消息
 * - LRANGE 获取完整对话
 * - LTRIM 保留最后 N 条（默认 15 条）
 * - EXPIRE 设置会话过期时间
 *
 * @author Albert
 * @since 2025-07-10 21:59:08
 */
@Component
public class RedisChatMemoryStore implements ChatMemoryStore {

    private static final int MAX_MESSAGES = 15;

    private final RedisTemplate<String, String> redisTemplate;

    public RedisChatMemoryStore(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    private String key(Object memoryId) {
        return RedisConstant.CHAT_MEMORY_PREFIX + memoryId;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String redisKey = key(memoryId);
        List<String> list = redisTemplate.opsForList().range(redisKey, 0, -1);
        if (CollUtil.isEmpty(list)) {
            return Collections.emptyList();
        }
        return list.stream()
                .map(ChatMessageDeserializer::messageFromJson)
                .collect(Collectors.toList());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String redisKey = key(memoryId);

        if (CollUtil.isEmpty(messages)) {
            return;
        }

        // 获取 Redis 当前已存储的消息数
        Long currentSize = redisTemplate.opsForList().size(redisKey);
        int alreadyStored = currentSize != null ? currentSize.intValue() : 0;

        // 只追加新消息
        if (messages.size() > alreadyStored) {
            List<String> newJsonList = messages.subList(alreadyStored, messages.size())
                    .stream()
                    .map(ChatMessageSerializer::messageToJson)
                    .collect(Collectors.toList());

            redisTemplate.opsForList().rightPushAll(redisKey, newJsonList);
        }

        // 只保留最近 MAX_MESSAGES 条
        redisTemplate.opsForList().trim(redisKey, -MAX_MESSAGES, -1);

        // 刷新过期时间
        redisTemplate.expire(redisKey, RedisConstant.CHAT_MEMORY_EXPIRE_TIME.getSeconds(), TimeUnit.SECONDS);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        redisTemplate.delete(key(memoryId));
    }
}
